import pandas as pd
from scipy.io import arff
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.impute import SimpleImputer
from sklearn.svm import SVC
from sklearn.metrics import classification_report, confusion_matrix, roc_curve, auc
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from util.LoadData import load_data
from charts.Draw import *

# 读取 .arff 文件并转换为 Pandas DataFrame
# def load_data(file_path):
#     # data 是数据部分，meta是包含数据集的元信息（比如属性名、数据类型等）
#     data, meta = arff.loadarff(file_path)
#     # 将 ARFF 数据转换为 Pandas DataFrame，以便进一步进行数据处理
#     df = pd.DataFrame(data)
#     # 处理数据类型， applymap 表示用于对DataFrame的每个元素应用一个函数，由于 ARFF 文件中可能存在字节类型数据，这里将字节数据解码为字符串。
#     df = df.applymap(lambda x: x.decode() if isinstance(x, bytes) else x)
#     return df

# 预处理数据，包括处理缺失值和转换标签
def preprocess_data(df):
    # 检查 'Defective' 列是否存在
    if 'Defective' not in df.columns:
      raise ValueError("DataFrame does not contain 'Defective' column.")
    imputer = SimpleImputer(strategy='mean') # 或选择 'median' 或 'most_frequent'
    df_imputed = pd.DataFrame(imputer.fit_transform(df.select_dtypes(include=['number'])), columns=df.select_dtypes(include=['number']).columns)
    #如果值是 'Y'，则映射为 1。如果值是 'N'，则映射为 0。 这样，'Defective' 列中的值就被转换为二进制格式，其中 1 表示有缺陷（Defective），0 表示无缺陷（Not Defective）
    df_imputed['Defective'] = df['Defective'].apply(lambda x: 1 if x == 'Y' else 0)
    return df_imputed

# 划分数据集为特征和标签，并进行标准化和数据集划分
def split_data(df):
    # 特征和标签
    # X 包含所有特征（即除去标签列的所有列）
    # y 是标签列（即 Defective）
    X = df.drop('Defective', axis=1)
    y = df['Defective']
    # 数据标准化
    # StandardScaler 用于标准化特征，使其具有均值为 0 和标准差为 1
    # fit_transform 对数据进行标准化
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    # 划分数据集
    # train_test_split 用于将数据集划分为训练集和测试集
    # test_size=0.3 表示 30% 的数据用于测试集，其余用于训练集
    # random_state=42 确保每次运行代码时划分结果一致。
    X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.3, random_state=42)
    return X_train, X_test, y_train, y_test

# 训练 SVM 模型
def train_svm(X_train, y_train):
    model = SVC(kernel='linear', probability=True)
    model.fit(X_train, y_train)
    return model

# 使用训练好的模型进行预测并评估模型性能
def predict_and_evaluate(model, X_test, y_test):
    # predict 方法生成测试数据的预测标签。
    # predict_proba 生成测试数据的预测概率，其中 [:, 1] 表示提取属于类别 1 （Y，不存在问题）的概率
    y_pred = model.predict(X_test)
    y_pred_prob = model.predict_proba(X_test)[:, 1]
    # classification_report 计算分类模型的精确度、召回率、F1 分数等指标，并以字典形式返回
    report = classification_report(y_test, y_pred, output_dict=True)
    # confusion_matrix 计算混淆矩阵，用于显示真实标签和预测标签的对比
    conf_matrix = confusion_matrix(y_test, y_pred)
    return y_pred, y_pred_prob, report, conf_matrix

# 绘制分类指标的直方图
def plot_metrics(metrics, title, filename):
    color = 'skyblue' if '0' in title else 'lightgreen'
    plt.figure(figsize=(6, 6))
    plt.bar(metrics.keys(), metrics.values(), color=color)
    plt.title(title)
    plt.ylabel('Score')
    plt.tight_layout()
    plt.savefig(filename)
    plt.show()

# 绘制 ROC 曲线和计算 AUC 值
def plot_roc_curve(y_test, y_pred_prob, filename):
    fpr, tpr, _ = roc_curve(y_test, y_pred_prob)
    roc_auc = auc(fpr, tpr)
    plt.figure(figsize=(8, 6))
    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (area = {roc_auc:.2f})')
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic (ROC)')
    plt.legend(loc='lower right')
    plt.tight_layout()
    plt.savefig(filename)
    plt.show()

# 绘制混淆矩阵的热图
def plot_confusion_matrix(conf_matrix, filename):
    plt.figure(figsize=(8, 6))
    sns.heatmap(conf_matrix, annot=True, fmt='d', cmap='Blues',
                xticklabels=['No Defect', 'Defect'],
                yticklabels=['No Defect', 'Defect'])
    plt.xlabel('Predicted')
    plt.ylabel('Actual')
    plt.title('Confusion Matrix')
    plt.savefig(filename)
    plt.show()

# 主函数，用于执行整个流程
def main():
    file_path = r'D:\EDAG下载\NASADefectDataset-master\OriginalData\MDP\CM1.arff'
    df = load_data(file_path)
    df_imputed = preprocess_data(df)
    X_train, X_test, y_train, y_test = split_data(df_imputed)
    model = train_svm(X_train, y_train)
    y_pred, y_pred_prob, report, conf_matrix = predict_and_evaluate(model, X_test, y_test)

    # 绘制评估指标
    metrics_0 = {
        'Accuracy': report['accuracy'],
        'Precision_0': report['0']['precision'],
        'Recall_0': report['0']['recall'],
        'F1_Score_0': report['0']['f1-score']
    }
    plot_metrics(metrics_0, 'Classification Metrics for Class 0 (Defect)', 'picture/SVM/classification_metrics_0.png')

    metrics_1 = {
        'Accuracy': report['accuracy'],
        'Precision_1': report['1']['precision'],
        'Recall_1': report['1']['recall'],
        'F1_Score_1': report['1']['f1-score']
    }
    plot_metrics(metrics_1, 'Classification Metrics for Class 1 (NO Defect)', 'picture/SVM/classification_metrics_1.png')

    # 绘制 ROC 曲线
    plot_roc_curve(y_test, y_pred_prob, 'picture/SVM/roc_curve_and_auc.png')

    # 绘制混淆矩阵
    plot_confusion_matrix(conf_matrix, 'picture/SVM/Confusion_Matrix.png')

# 程序入口点
if __name__ == '__main__':
    main()